
# from socket import *  
import socket
import time
import re
import select
import ipmsg_protocol  as pt 
from down_file_depack import *
from threading import Thread


class usr_local:
    def __init__(self, user='local', host='localhost'):
        self.host = host
        self.usr_name = user 
        self.pt_version = 1
        self.msg_id = int(time.time())

    def pack_msg(self, msg_opt, msg_data='',msg_mode=0):
        self.msg_id += 1
        return '%d:%d:%s:%s:%d:%s'%(self.pt_version,self.msg_id, self.usr_name,self.host,msg_opt|msg_mode,str(msg_data)) 

    def get_msid(self):
        self.msg_id += 1
        return self.msg_id

class ipmsg(Thread):
    def __init__(self, wg):
        Thread.__init__(self)
        self.wg = wg 
        self.o_usr = usr_local()
        self._exit = False 

        self.file_id = [0,0,0,0,0,0,0,0,0,0] 
        self.file_names = {}
        self.file_wid = [0,0,0,0,0,0,0,0,0,0] 
        self.file_widget = [0,0,0,0,0,0,0,0,0,0] 

        self.connections = {}
        self.addresses = {}

    def get_mode(self, command):
        return command & 0x000000ff

    def get_opt(self, command):
        return command & 0xffffff00 

    def msg_depack(self, data):
        try:
            data = data.decode('gbk')
        except Exception as e:
            print(e)
            print(data)
        ts = data.split('::')
        msg = ['']
        for t in ts:
            ds = t.split(':') 
            for i,d in enumerate(ds):
                if i == 0:
                    msg[-1] += ':'+d
                else:
                    msg.append(d)
        s_msg = {}
        lmsg = len(msg)
        if lmsg >3:
            s_msg['ver'] = msg[0]
            s_msg['id'] = int(msg[1])
            s_msg['usr_name'] = msg[2]
            s_msg['usr_host'] = msg[3]
        if lmsg >4:
            t = int(msg[4])
            s_msg['opt'] = self.get_opt(t)
            s_msg['mode'] = self.get_mode(t)
        if lmsg >5:
            s_msg['data'] = msg[5:]
        return s_msg

    def msg_encode(self, data):
        return data.encode('iso8859-1')

    def msg_ept_pack(self, data):
        ds = ''
        for d in data:
            d.replace(':','::')
            if ds != '':
                ds += ':'
            ds += d 
        return (ds) #.encode('iso8859-1') self.msg_encode
    #['\x000', 'msdfmap.ini', '57d', '4a301960', '1', '\x001', 'PFRO.log', '12c', '5fa5f818', '1', '\x00']
    def file_info_depack(self, data):
        file_infos = [] 
        file_attr = ['id','name','size','mtime','attr']
        fattr_l = len(file_attr)
        fattr_i = 0

        file_info = None
        for d in data:
            if d[0]=='\x00':
                if file_info is None:
                    pass
                elif len(file_info) >= fattr_l:
                    file_info['off'] = 0
                    file_infos.append(file_info)
                else:
                    print('error:',data,'  |->|  ', file_infos)
                    return []
                file_info = {}
                fattr_i = 0 

                file_info[file_attr[fattr_i]] = d[1:] 
                fattr_i += 1
            else:
                file_info[file_attr[fattr_i]] = d
                fattr_i += 1
        return file_infos 

    def down_file(self, sock, down_info):
        recv_d = sock.recv(1000000) # 65535 
        if len(recv_d) <= 0:
            return False
        dx =down_info 
        msg_item = dx['msg_item']
        save_dir = msg_item.file_info['save_dir'] 
        print(msg_item.file_info , dx) 
        if 'down_size' not in dx: dx['down_size'] = 0
        # down_size = 0 
        down_finish = False
        while len(recv_d) > 0: #and msg_item.dwn_btn_stu: 
            # down_size += len(recv_d) 
            dx['down_size'] += len(recv_d)
            print('down_size:',dx['down_size'])
            msg_item.pbar.setValue(int(dx['down_size']/dx['total_size']*100))
            if dx['down_size'] >= dx['total_size']:
                down_finish = True 
            # msg_item.dwn_btn.setText(msg_item.dwn_btn_txt[msg_item.dwn_btn_stu]) #setValue(int(down_size/dx['total_size']*100))
            # img_item.file_info. file_info
            data_off = 0
            dx['last_data'] += recv_d
            data = dx['last_data']
            dl = len(data) 
            # print(dx)
            wl = min(dl, dx['file_size']-dx['file_off']) 
            if wl > 0 and dx['file_fp'] is not None:
                dx['file_fp'].write(data[data_off:wl])
                data_off += wl 
                dx['file_off'] += wl  

            if dx['file_size']-dx['file_off'] == 0:
                if dx['file_fp'] is not None:
                    dx['file_fp'].close()
                    dx['file_fp'] = None  
                    dx['file_size']=dx['file_off'] = 0 
            while dl > data_off:
                finfo = file_data_depack(data[data_off:]) 
                # print('down:',data[data_off:100].decode('ISO-8859-1')) 
                print('down:', end='')
                if finfo is None: # 数据错误，直接关掉传输管道
                    # self.epoll.modify(fd, select.EPOLLHUP)
                    # break
                    return False
                elif len(finfo) == 0: # 数据长度不够，等待下次传输
                    break 
                # 目录文件信息 
                if finfo['fattr']== pt.IPMSG_FILE_DIR:
                    dx['dir'].append(finfo['fname'])
                    dx['dir_size'].append(finfo['fsize'])
                    dir_path = os.path.join(save_dir,'/'.join(dx['dir']))
                    if not os.path.exists(dir_path):
                        os.makedirs(dir_path) 
                    data_off += finfo['hsize']
                # 目录返回信息
                elif finfo['fattr']== pt.IPMSG_FILE_RETPARENT:
                    dx['dir'].pop(-1)
                    dx['dir_size'].pop(-1)
                    data_off += finfo['hsize'] 
                # 其他文件信息
                else:
                    data_off += finfo['hsize'] 
                    dx['file'] = finfo['fname'] 
                    dx['file_size'] = finfo['fsize'] 
                    dx['file_off'] = 0
                    fpath = os.path.join(save_dir,'/'.join(dx['dir']), dx['file']) 
                    dx['file_fp'] = open(fpath,'wb') 

                    wl = min(dl-data_off, dx['file_size']-dx['file_off'])
                    if wl > 0 :
                        dx['file_fp'].write(data[data_off:data_off+wl])
                        data_off += wl 
                        dx['file_off'] += wl 

                    if dx['file_size']-dx['file_off'] == 0:
                        if dx['file_fp'] is not None:
                            dx['file_fp'].close()
                            dx['file_fp'] = None 
                    data_off += finfo['fsize'] 
            dx['last_data'] = data[data_off:] # 将剩余数据保存起来 
            try: 
                recv_d = sock.recv(1000000) 
                # data = sock.recv(1024)
            except BlockingIOError as err:
                break 
        if down_finish:
            print('3传输完成！！')
            msg_item.set_btn('已完成')
            self.epoll.modify(sock.fileno(), select.EPOLLHUP) 
        return True

    def down_file_pack(self, msg_id, file_info):
        return self.msg_ept_pack(['%x'%(msg_id),'%s'%file_info['id'],'%x'%file_info['off']]) 

    def send_msg(self,address, msg):
        self.udp_sd.sendto(self.o_usr.pack_msg(pt.IPMSG_SENDMSG|pt.IPMSG_ABSENCEOPT,msg.replace(':','::')).encode('gbk'),address) 
    
    def send_down_file_msg(self, address, msg_item):
        msg_data = msg_item.file_info
        file_infos = self.file_info_depack(msg_data['data']) 
        if address not in self.message_queues:
            self.message_queues[address] = msg_item 
        usr_data = self.message_queues[address]
        if 'file' not in usr_data:
            usr_data['file'] = {} 
        usr_data['file'][msg_data['id']] = file_infos
        # # 下载 
        for f in file_infos:
            sd_data = self.o_usr.pack_msg(pt.IPMSG_GETFILEDATA,self.down_file_pack(msg_data['id'], f))
            tcp_sub = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            tcp_sub.connect(address) 
            tcp_sub.send(sd_data.encode('gbk')) 
            tcp_sub.setblocking(False)
            print('add download epoll!')
            self.epoll.register(tcp_sub.fileno(), select.EPOLLIN|select.EPOLLHUP) 
            self.message_queues[tcp_sub.fileno()] = {'file_id':f['id'],'total_size':int(f['size'],base=16),'ntry_cnt':0,'msg_item':msg_item,'dir':[], 'dir_size':[],  'file':'', 'file_fp':None, 'file_size':0, 'file_off':0, 'last_data':b''} #'dir_off':[],
            self.fd_to_socket[tcp_sub.fileno()] = tcp_sub 

    def get_size(self, fpath = '.'):
        if not os.path.isdir(fpath):
            return os.path.getsize(fpath)
        total_size = 0
        for dirpath, _, filenames in os.walk(fpath):
            folder_size = 0
            for f in filenames:
                fp = os.path.join(dirpath, f)
                if not os.path.islink(fp):
                    folder_size += os.path.getsize(fp)
            total_size += folder_size
        return total_size 

    def get_file_type(self, f):
        return  pt.IPMSG_FILE_DIR  if os.path.isdir(f) else pt.IPMSG_FILE_REGULAR 

    def upfile_head(self, fpath):
        fn = os.path.basename(fpath)
        hd = [fn]
        if os.path.isdir(fpath):
            hd += ['%d'%self.get_size(fpath)]
        else:
            hd += ['%09x'%self.get_size(fpath)]
        hd += ['%x'%self.get_file_type(fpath),
        '%x=%x'%(pt.IPMSG_FILE_MTIME,int(os.path.getmtime(fpath))), 
        '%x=%x'%(pt.IPMSG_FILE_CREATETIME,int(os.path.getctime(fpath)))] 
        file_hd = ':'.join(hd)+":"
        print('upfile:','%04x:'%(len(file_hd)+5)+file_hd)
        return '%04x:'%(len(file_hd.encode('gbk'))+5)+file_hd 

    def fdown_cancel(self, msg_item):

        # if  msg_item.dwn_btn_stu == 2:
        if msg_item.sock is not None:
            self.udp_sd.sendto(self.o_usr.pack_msg(pt.IDS_EXCEPTIONLOG, str(msg_item.file_info['id'])+dx['file_id']+'\x00'+self.o_usr.usr_name+'\x00').encode('gbk'),msg_item.file_info['usr_addr'])
            self.epoll.modify(msg_item.sock.fileno(), select.EPOLLHUP) 
        # msg_item.dwn_btn.setText)#(msg_item.dwn_btn_txt[msg_item.dwn_btn_stu]
        msg_item.set_btn('已取消')
        msg_item.sock = None
        # msg_item.dwn_btn.setEnabled(False)
        # print('cancel') 
            # return 
        print('fdown cancel')
    
    def upload_send_cancel(self, msg_item):
        print(msg_item.sock,self.wg.cur_usr) #fs_id 
        for fid in msg_item.fs_id:
            self.udp_sd.sendto(self.o_usr.pack_msg(pt.IDS_UNDO, '\x00'+str(self.file_wid[fid])+str(fid)+'\x00').encode('gbk'),msg_item.file_info['usr_addr'])
        msg_item.set_btn('已取消')
        # print('cancel upload')

    def upfile_dirret_head(self):
        # fn = os.path.basename(fpath)
        hd = ['.','%x'%0,'%x'%pt.IPMSG_FILE_RETPARENT]
        file_hd = ':'.join(hd)+":"
        print('upfile ret dir:','%04x:'%(len(file_hd)+5)+file_hd)
        return '%04x:'%(len(file_hd.encode('gbk'))+5)+file_hd 

    def upload_file(self, sock):
        recv_d = sock.recv(10000) #spath 
        dc_data = self.msg_depack(recv_d) 
        fid = int(dc_data['data'][1],base=16)
        print('dc_data:',dc_data)
        # print('upload:',recv_d[:100], dc_data)
        # print(self.file_names[fid])
        spath = self.file_names[fid]
        # return
        old_path = ''

        if not os.path.isdir(spath):
            sock.send(self.upfile_head(spath).encode('gbk'))
            with open(spath,'rb') as fr:
                print('send file')
                sock.send(fr.read()) 
            print('2传输完成！！')
            self.file_widget[int(dc_data['data'][1])].pbar.setValue(100)
            self.file_widget[int(dc_data['data'][1])].set_btn('已完成')
            self.epoll.modify(sock.fileno(), select.EPOLLHUP) 
            # print('关闭套接字')
            return 
        
        for dirpath, _, filenames in os.walk(spath):
            print(dirpath, filenames)
            mlen = min(len(old_path), len(dirpath))
            if mlen != 0 and old_path[:mlen] != dirpath[:mlen]:
                self.upfile_dirret_head()
                old_path = dirpath
            sock.send(self.upfile_head(dirpath).encode('gbk'))
            print(filenames)
            for f in filenames:
                fp = os.path.join(dirpath, f) 
                sock.send(self.upfile_head(fp).encode('gbk'))
                with open(fp,'rb') as fr:
                    print('send file')
                    sock.send(fr.read()) 
        print('1传输完成！！')
        self.file_widget[int(dc_data['data'][1])].pbar.setValue(100)
        self.file_widget[int(dc_data['data'][1])].set_btn('已完成')
        # msg_item.set_btn('已完成')
        self.epoll.modify(sock.fileno(), select.EPOLLHUP)
                    
    def send_upload_file_info(self, sock, fps, recv_usr, widget):
        usd_id = []
        hd = [] 

        for fp in fps: 
            fid = self.file_id.index(0) 
            self.file_id[fid]=1 
            usd_id.append(fid)
            if fid < 0:
                return [False,[]]
                # return [False] 
            self.file_names[fid] = fp 
            if hd != []:
                hd += ['\a']
            hd += ['%x'%fid, os.path.basename(fp), '%x'%self.get_size(fp), '%x'%int(os.path.getmtime(fp)), '%x'%self.get_file_type(fp)] 

        fs_data = ':'.join(hd)+':'
        fs_data = '\0'+fs_data+'\0' 
        sock.sendto(self.o_usr.pack_msg(pt.IPMSG_SENDMSG|pt.IPMSG_FILEATTACHOPT|pt.IPMSG_ABSENCEOPT,fs_data).encode('gbk'),recv_usr)
        for uid in usd_id:
            self.file_wid[uid] = self.o_usr.msg_id 
            self.file_widget[uid] = widget 

        return [True, usd_id]

    def reset_file_id(fid):
        self.file_id[fid] = 0

    def run(self):
        data = self.o_usr.pack_msg(pt.IPMSG_BR_ENTRY)

        self.udp_sd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_sd.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        self.udp_sd.bind(('0.0.0.0', pt.IPMSG_DEFAULT_PORT)) 
        
        timeout = 0.5
        self.message_queues = {}
        self.epoll = select.epoll()

        self.epoll.register(self.udp_sd.fileno(), select.EPOLLIN|select.EPOLLHUP)
        self.fd_to_socket = {self.udp_sd.fileno():self.udp_sd,}

        self.tcp_sd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tcp_sd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.tcp_sd.bind(('0.0.0.0', pt.IPMSG_DEFAULT_PORT))
        self.tcp_sd.listen(128)

        self.epoll.register(self.tcp_sd.fileno(), select.EPOLLIN|select.EPOLLHUP)
        self.fd_to_socket[self.tcp_sd.fileno()] = self.tcp_sd

        self.udp_sd.sendto(data.encode('utf-8'), (pt.IPMSG_LIMITED_BROADCAST, pt.IPMSG_DEFAULT_PORT))

        while not self._exit:
            # print("\r等待活动连接......",end='')
            events = self.epoll.poll(timeout)
            if not events:
                # print ("\repoll超时无活动连接，重新轮询......",end='')
                continue
            # print ("\r有" , len(events), "个新事件，开始处理......",end='')
        
            for fd, event in events:
                if fd not in self.fd_to_socket:
                    continue
                sock = self.fd_to_socket[fd]
                if sock == self.udp_sd:
                    data, address = self.udp_sd.recvfrom(65535) 
                    print(data)
                    dc_data = self.msg_depack(data) 
                    dc_data['usr_addr'] = address
                    # 接受到 广播登记 消息 
                    if dc_data['mode'] == pt.IPMSG_BR_ENTRY:
                        self.message_queues[address] = dc_data
                        self.wg.add_usr.emit(self.message_queues[address])
                        # print('add:',address)
                        #发送 登记应答 消息 
                        sock.sendto(self.o_usr.pack_msg(pt.IPMSG_ANSENTRY,self.o_usr.usr_name).encode('gbk'),address) 
                    # 接受到 登记应答 消息 
                    elif dc_data['mode'] == pt.IPMSG_ANSENTRY:
                        udp_sub = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                        udp_sub.connect(address)
                        self.epoll.register(udp_sub.fileno(), select.EPOLLIN)
                        self.message_queues[address] = dc_data 
                        # print('add:',address)
                        self.wg.add_usr.emit(self.message_queues[address])
                        self.fd_to_socket[udp_sub.fileno()] = udp_sub 
                    # 接受到 通信信息 信息
                    elif dc_data['mode'] == pt.IPMSG_SENDMSG:
                        # 消息应答
                        if  dc_data['opt']&pt.IPMSG_ABSENCEOPT:
                            sock.sendto(self.o_usr.pack_msg(pt.IPMSG_RECVMSG,dc_data['id']).encode('gbk'),address)
                        else:
                            sock.sendto(self.o_usr.pack_msg(pt.IPMSG_RECVMSG,self.o_usr.usr_name).encode('gbk'),address) 
                        
                        self.wg.add_msg.emit(self.message_queues[address], dc_data)
                    elif dc_data['mode'] == pt.IPMSG_RECVMSG:
                        print(dc_data)
                    else:
                        print('unkonw[%s:%d]:'%(address), dc_data['opt']|dc_data['mode']) 
                        print(dc_data) 
                if sock == self.tcp_sd:
                    if fd not in self.message_queues:

                        conn,addr=sock.accept()
    
                        print('有新的客户端到来%s'%str(addr))
    
                        # 将 conn 和 addr 信息分别保存起来
                        self.fd_to_socket[conn.fileno()] = conn 
                        self.addresses[conn.fileno()] = addr 
                        self.connections[conn.fileno()] = addr 

                        self.epoll.register(conn.fileno(), select.EPOLLIN | select.EPOLLET) 
                    # data = sock.recv(1000000) 
                    # print('client:',data)       

                elif event & select.EPOLLHUP:
                    self.epoll.unregister(fd)
                    if fd in self.fd_to_socket:
                        del self.fd_to_socket[fd]
                    # del self.message_queues[fd]
                    print('close:',sock.fileno())
                    sock.close()
                #可读事件
                elif event & select.EPOLLIN: 

                    if fd in self.connections:
                        self.upload_file(sock)
                        print("上传成功！！")

                    elif fd in self.message_queues:
                        # print(fd, fd in self.message_queues)
                        dsuc = self.down_file(sock, self.message_queues[fd])
                        if not dsuc:
                            self.epoll.modify(fd, select.EPOLLHUP) 
                #可写事件 
                elif event & select.EPOLLOUT:
                    sock.send('abc'.encode('gbk'))
                else : 
                    print(fd, 'can\'t parse')
    def exit(self):
        self._exit = True

if __name__ == "__main__":
    ipmsg_t()